#!/usr/bin/python
# -*- coding: utf-8 -*-
################################################################################
# Licensed to the OpenAirInterface (OAI) Software Alliance under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The OpenAirInterface Software Alliance licenses this file to You under 
# the Apache License, Version 2.0  (the "License"); you may not use this file
# except in compliance with the License.  
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#-------------------------------------------------------------------------------
# For more information about the OpenAirInterface (OAI) Software Alliance:
#      contact@openairinterface.org
################################################################################
import sys
import subprocess
import re
import socket
import datetime
from datetime import date
import os, errno
import argparse
import tempfile
from lxml import etree
from xml.dom.minidom import parse, parseString

#####################
# program arguments
#####################
parser = argparse.ArgumentParser()
parser.add_argument("--pcap_file", "-p", type=str,help="input pcap file to be translated")
args = parser.parse_args()

#####################
# get xml document from pcap
#####################
orig_pcap_file_name = args.pcap_file.strip()

## NO SCTP_CID_SACK(3), SCTP_CID_HEARTBEAT, SCTP_CID_HEARTBEAT_ACK, SCTP_CID_COOKIE_ECHO, SCTP_CID_COOKIE_ACK
#orig_pdml_string = subprocess.check_output(["tshark", '-2', '-F', 'pcapng' , '-w', orig_pcap_file_name + ".SCTP_Filtered.pcapng", '-t', 'r', '-r', orig_pcap_file_name, '-R', 'sctp.chunk_type!=3 && sctp.chunk_type!=4 && sctp.chunk_type!=5 && sctp.chunk_type!=10 && sctp.chunk_type!=11'])
# !!! Bug in tshark: it zeroes all timestamps !!!


orig_pdml_string = subprocess.check_output(["tshark", '-2', '-T', 'pdml', '-r', orig_pcap_file_name])

orig_dom = parseString(orig_pdml_string)

#####################
# filtering unwanted packets
#####################
#cases = orig_etree.findall(".//proto[@name='sctp']")
packets = orig_dom.getElementsByTagName("packet")
for packet in packets:
    found_sctp = False
    found_s1ap = False
    sctp_node = None
    frame_node = None
    ip_node = None
    protos = packet.getElementsByTagName("proto")
    for proto in protos:
        attrs   = proto.attributes
        urlnode = attrs['name']
        if urlnode.nodeValue == 'frame':
            frame_node = proto
        elif urlnode.nodeValue == 'ip':
            frame_node.appendChild(proto)
            ip_node = proto
        elif urlnode.nodeValue == 'sctp':
            fields = proto.getElementsByTagName("field")
            for field in fields:
                if found_sctp == True:
                    continue
                attrs_field = field.attributes
                attr_field = attrs_field['name']
                if attr_field.nodeValue == 'sctp.chunk_type':
                    attr_field = attrs_field['value']
                    # 00 = DATA CHUNK, 01 = INIT CHUNK, 02 = INIT_ACK_CHUNK
                    if (attr_field.nodeValue <= '02') or (attr_field.nodeValue == '07') or (attr_field.nodeValue == '08')  or (attr_field.nodeValue == '0e') : 
                        found_sctp = True
            ip_node.appendChild(proto)
            sctp_node = proto
        elif urlnode.nodeValue == 's1ap':
            found_s1ap = True
            sctp_node.appendChild(proto)
        elif urlnode.nodeValue == 'geninfo':
            packet.removeChild(proto)
        elif urlnode.nodeValue == 'eth':
            packet.removeChild(proto)
    if found_sctp == False:
        # hopefully it seems to work (remove iterated packet)
        packet.parentNode.removeChild(packet)

#####################
# dom to xml string
#####################
filtered_pdml_string = orig_dom.toprettyxml(indent="  ")
cleaned_pdml_string = ""
#####################
# remove blank lines in xml string
#####################
lines = filtered_pdml_string.splitlines()
for line in lines:
    if len(line.strip()):
        if line[:-1]:
            cleaned_pdml_string += line + '\r\n'
#print "'%s'" %  cleaned_pdml_string 
#####################
# write pdml string to pdml file
#####################
out_pdml_file_name = os.path.dirname(orig_pcap_file_name) + '/' + os.path.splitext(os.path.basename(orig_pcap_file_name))[0] + '.pdml'
out_file = open(out_pdml_file_name, "w")
out_file.write(cleaned_pdml_string)
out_file.close()

############################################################
# DECEIVING HTML BONUS: DO NOT SEEM TO WORK CORRECTLY IN FIREFOX
# DID NOT INVESTIGATE
#####################
# write xml string to html file
#####################
xsl_root = etree.fromstring(open('/usr/share/wireshark/pdml2html.xsl').read())
transform = etree.XSLT(xsl_root)
xml_root = etree.fromstring(cleaned_pdml_string)
trans_root = transform(xml_root)
filtered_html_string = etree.tostring(trans_root)
#####################
# write html string to html file
#####################
out_html_file_name = os.path.dirname(orig_pcap_file_name) + '/' +os.path.splitext(os.path.basename(orig_pcap_file_name))[0] + '.html'
out_file = open(out_html_file_name, "w")
out_file.write(filtered_html_string)
out_file.close()
